from dataclasses import dataclass
from enum import Enum
from typing import Dict, List, Optional, Tuple, Union


@dataclass(frozen=True, eq=True)
class Entity:
    """Entity dataclass
    :param name(str): entity name
    :param type(str): entity type
    """

    name: str
    type: str = "entity"

    @staticmethod
    def from_json(data):
        return Entity(
            name=data["name"],
            type=(
                str.lower(data["type"])
                if data["type"] not in ["NA", "na", None, ""]
                else "entity"
            ),
        )

    @staticmethod
    def from_tuple(data):
        return Entity(name=data[0], type=data[1])

    @staticmethod
    def to_str(data):
        return data.name + " " + data.type


@dataclass(frozen=True, eq=True)
class Relation:
    """Relation dataclass
    :param head(str): head entity
    :param tail(str): tail entity
    :param type(str): relation type
    """

    head: str
    tail: str
    type: str

    @staticmethod
    def from_json(data):
        head = data["head"]
        tail = data["tail"]
        return Relation(
            head=head["name"],
            tail=tail["name"],
            type=data["type"].lower(),
        )

    @staticmethod
    def from_tuple(data):
        return Relation(
            head=data[0],
            tail=data[1],
            type=data[2],
        )

    def to_tuple(self):
        return (self.head, self.tail, self.type)

    def __eq__(self, value: object) -> bool:
        if not isinstance(value, Relation):
            return False
        return (
            self.head == value.head
            and self.tail == value.tail
            and self.type == value.type
        )

    def __iter__(self):
        return iter((self.head, self.tail, self.type))


@dataclass()
class Instance:
    """
    Core Instance dataclass.
    :param context (str): input sentence and some conditions or question before answer
    :param instruction (str): instruction + alternatives
    :param options (str): options for QA model
    :param answer (optional): raw answer from QA model
    """

    context: str
    instruction: str
    answer: Optional[str]
    relations: Optional[List[Tuple[str]]]

    def __str__(self):

        result = []

        result.append(self.instruction)

        result.append(self.context)

        if self.answer is not None:
            result.append(self.answer)

        return " ".join(result)


@dataclass
class Example:
    """
    Example dataclass
    :param sentence (str): input sentence
    :param relations (str): relation answer
    """

    sentence: str
    pefix: Optional[str]
    answer: str

    def __str__(self) -> str:
        if self.prefix is not None:
            return self.sentence + " " + self.prefix + " " + self.answer
        return self.sentence + " " + self.answer


class DatasetType(Enum):
    CONLL2004 = "conll2004"


class Preffix(Enum):
    CONTEXT = "Sentence: "
    INSTRUCTION = "Instruction: "
    OPTIONS = "Options: "
    ANSWER = "Answer: "
    CONDITIONS = "Conditions: "
    QUESTION = "Question: "
    EXAMPLES = "Examples: "


class TaskType(Enum):
    RE = "RE"
    RC_NO_ICL = "RC_NO_ICL"
    RC = "RC"
    TE = "TE"
    OE = "OE"
    SE = "SE"
    EPE = "EPE"
    TE_NO_OPTIONS = "TE_NO_OPTIONS"
    RC_W_EXP = "RC_W_EXP"
    NER = "NER"
    RE_STRICT = "RE_STRICT"


class DatasetField(Enum):
    SENTENCE = "sentence"
    RELATIONS = "relations"


class Language(Enum):
    EN = "en"


class AnswerType(Enum):
    TRIPLE = "triple"
    SVO = "svo"
    TEMPLATE = "template"


# 函数用于根据字符串值获取对应的枚举成员
def string_to_enum(string):
    for member in AnswerType:
        if member.value == string:
            return member
    raise ValueError(f"{string} is not a valid value for AnswerType")


class ContextFormat(Enum):
    CONDITIONS = "conditions"
    TAG = "tag"


def string_list_to_task_type(string_list):
    task_types = []
    for string in string_list:
        for member in TaskType:
            if member.value == string:
                task_types.append(member)
                break
    return task_types
